import os
import shutil
import tempfile
import subprocess

from src.utils.process import ProcessManager
from src.ui.i18n import _


class InterruptedError(Exception):
    pass

class VocalRemover:
    """
    Separates vocals and background music from an audio track using the Demucs AI model.
    """
    def __init__(self, process_manager: ProcessManager, progress_callback=None, work_dir=None, is_stopped_func=None):
        self.process_manager = process_manager
        self.progress_callback = progress_callback
        self.work_dir = work_dir or tempfile.mkdtemp()
        self.is_stopped = is_stopped_func or (lambda: False)

    def _execute_command(self, command):
        if self.is_stopped():
            raise InterruptedError(_("Task was stopped by the user."))
        
        # The `run` method is blocking and will raise an exception on failure.
        # It also handles logging via the provided callback.
        self.process_manager.run(command, log_callback=self.progress_callback)

    def separate(self, audio_path):
        """
        Separates vocals and accompaniment from a given audio file.
        :param audio_path: Path to the input audio file.
        :return: Path to the accompaniment file.
        """
        if self.progress_callback:
            self.progress_callback(_("Processing audio with Demucs (first run may download models)..."))

        model_name = "htdemucs_ft"
        output_dir = self.work_dir
        
        command = [
            "python3", "-m", "demucs",
            "--two-stems=vocals",
            "-n", model_name,
            "-o", output_dir,
            audio_path
        ]

        try:
            self._execute_command(command)
        except Exception as e:
            raise RuntimeError(_("Demucs execution failed: {}").format(e))

        # Check if the output file exists
        # Demucs saves the result in a subdirectory named after the model
        accompaniment_path = os.path.join(output_dir, model_name, f"{os.path.splitext(os.path.basename(audio_path))[0]}.wav")
        if not os.path.exists(accompaniment_path):
            raise RuntimeError(_("Demucs failed to generate the accompaniment file. Expected path: {}").format(accompaniment_path))

        # Move the file to the root of the work directory for easier reference
        final_accompaniment_path = os.path.join(self.work_dir, "accompaniment.wav")
        shutil.move(accompaniment_path, final_accompaniment_path)
        
        # Clean up the model subdirectory created by demucs
        shutil.rmtree(os.path.join(self.work_dir, model_name))

        if self.progress_callback:
            self.progress_callback(_("Vocal separation successful!"))

        return final_accompaniment_path
